/*
    Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/// \file
/// Extra mcrx functions. Functions to  read and write lost energy and 
/// random number generator states, calculate integrated SED and
/// dust-attenuated images. \ingroup mcrx

// $Id$

#include "mcrx.h"

#include "blitz-fits.h"
#include "biniostream.h"
#include "fits-utilities.h"
#include "interpolatort.h"
#include "constants.h"
#include "misc.h"
#include "vecops.h"
#include "imops.h"
#include "mono_poly_abstract.h"
#include "boost/lambda/lambda.hpp"
#include "boost/lambda/bind.hpp"

using namespace blitz;
using namespace CCfits;
using namespace boost::lambda;

/** Calculates the integrated SED by adding all the pixels in the
    cameras. This function is ridiculously long, I don't know why it
    has to be so complicated... */
bool mcrx::Mcrx::generate_integrated_SED ()
{
  if (terminator ()) return true;
  
  typedef interpolator< T_float , T_float , 1> T_interpolator;

  cout << "\n  *** ENTERING POSTPROCESSING STAGE ***\n" << endl;

  // check if we actually need to do this
  bool dump_sed= false;
  ofstream dump;
  if (p.defined("sed_file")) {
    dump_sed = true;
    dump.open (word_expand (p.getValue("sed_file" , string ())) [0].c_str ());
  }
  
  const string output_file_name = 
    word_expand (p.getValue("output_file", string ())) [0];
  auto_ptr<FITS> output;
  output.reset(new FITS (output_file_name, Read));

  // check if previously complete
  bool reentry = false;
  try {
    output->pHDU ().readKey("MCRX_PP1", reentry);
  }
  catch (...) {}
  if (reentry) {
    cout << "  Stage already complete" << endl;
    return false;
  }

  output.reset(0 );
  output.reset(new FITS (output_file_name, Write));
  ExtHDU& lambda_hdu = open_HDU (*output, "LAMBDA");
  Column& cl = lambda_hdu.column("lambda" );
  Column& cli = lambda_hdu.column("lambda_intensity" );
  units ["wavelength"] = cl.unit ();
  array_1 wavelengths;
  array_1 iwavelengths;
  read (cl, wavelengths);
  read (cli, iwavelengths);
  int n_lambda_int;
  lambda_hdu.readKey("NLAMBDA_INTENSITY", n_lambda_int);
  iwavelengths.resizeAndPreserve(n_lambda_int);

  // Check for terminator
  if (terminator ()) return true;

  int n_cameras;
  open_HDU (*output, "MCRX").readKey("N_CAMERA", n_cameras);
  // start by getting the source L_lambda, which requires reopening
  // the input file. this is done this early because we also need the
  // L_lambda unit
  array_1 L_lambda_total (wavelengths.size());
  {
    cout << "Reading integrated source SED" << endl;
    FITS input (word_expand (p.getValue("input_file", string ()))[0], Read);
    ExtHDU& data_hdu = open_HDU (input, "INTEGRATED_QUANTITIES");
    read (data_hdu.column("L_lambda"), L_lambda_total);
    assert (all (L_lambda_total == L_lambda_total));

    units ["L_lambda"] =data_hdu.column("L_lambda").unit();
    ExtHDU& pdata_hdu = open_HDU (input, "PARTICLEDATA");
    units ["luminosity"] =pdata_hdu.column("L_bol").unit();   
  }

  // create INTEGRATED_QUANTITIES HDU
  ExtHDU* iq_HDU;
  try {
    iq_HDU = &open_HDU (*output, "INTEGRATED_QUANTITIES");
  }
  catch (FITS::NoSuchHDU&) {
    // didn't exist, create it.
    iq_HDU = output->addTable("INTEGRATED_QUANTITIES", 0 );
    iq_HDU->writeComment("This HDU contains spatially integrated SED's of the various quantities");
    iq_HDU->writeComment("For the camera quantities, the luminosity is calculated from the camera flux by multiplying by 4*pi*d^2.  It's not a fundamental quantity.");
  }
  try {
    iq_HDU->column("L_lambda_absorbed");
  }
  catch (CCfits::Table::NoSuchColumn&) {
    // we haven't created the columns yet, do so
    iq_HDU->addColumn(Tdouble, "lambda", 1, units.get("wavelength"));
    iq_HDU->addColumn(Tdouble, "lambda_intensity", 1, units.get("wavelength"));
    iq_HDU->addColumn(Tdouble, "L_lambda", 1, units.get("L_lambda"));
    iq_HDU->addColumn(Tdouble,"L_lambda_absorbed", 1, units.get("L_lambda"));
    // also create the camera-specific columns, otherwise we will need
    // to create space in the FITS file later
    for (int j=0;j < n_cameras; ++j) {
      std::ostringstream ost;
      ost << j;
      const string camera_numstr=ost.str();	
      // only create columns if the corresponding camera exists
      try {
	open_HDU(*output, "CAMERA"+camera_numstr+"-NONSCATTER");
	iq_HDU->addColumn(Tdouble, "L_lambda_nonscatter" +camera_numstr, 1,
			  units.get("L_lambda"));
      }
      catch (CCfits::FitsError&) {}
      catch (CCfits::FITS::NoSuchHDU&) {}
      try {
	open_HDU(*output, "CAMERA"+camera_numstr);
	iq_HDU->addColumn(Tdouble, "L_lambda_scatter" +camera_numstr, 1, 
			  units.get("L_lambda"));
      }
      catch (CCfits::FitsError&) {}
      catch (CCfits::FITS::NoSuchHDU&) {}
      try {
	open_HDU(*output, "CAMERA"+camera_numstr+"-IR");
      iq_HDU->addColumn(Tdouble, "L_lambda_ir" +camera_numstr, 1, 
			units.get("L_lambda"));
      }
      catch (CCfits::FitsError&) {}
      catch (CCfits::FITS::NoSuchHDU&) {}
    }
  }

  write (iq_HDU->column("lambda"), wavelengths, 1 );
  write (iq_HDU->column("lambda_intensity"), iwavelengths, 1 );
  iq_HDU->addKey("nlambda_intensity",n_lambda_int,"");
  write (iq_HDU->column("L_lambda"), L_lambda_total, 1 );
  const T_float L_bol_grid = integrate_quantity (L_lambda_total,
						 wavelengths, true);
  iq_HDU->addKey("L_bol_grid", L_bol_grid, "[" + units.get("luminosity") +
		 "] Bolometric source luminosity");

  // Check for terminator
  if (terminator ()) return true;

  // Now deal with the absorbed luminosity
  array_1 deposition_sed (calculate_L_absorbed(*output));

  cout << "Bolometric luminosity in grid: " << L_bol_grid << " " <<
    units.get("luminosity") << endl;

  cout << "\nCalculating camera quantities" << endl;

  if (dump_sed) {
    // dump  total and absorbed SED 
    for (array_1::iterator i = wavelengths.begin(),
	   j =  L_lambda_total.begin(),
	   l = deposition_sed.begin();
	 i != wavelengths.end(); ++i, ++j, ++l) {
      dump  <<*i << '\t' <<*j  << '\t' <<*l 
	    <<'\n';
    } 
    dump << "\n\n";
  }
  
  // Check for terminator
  if (terminator ()) return true;

  // Loop over Cameras

  int i = 0;
  // See if we should resume at a specific camera
  try {
    output->pHDU().readKey("MCRX_PP1_RESUME", i );
    output->pHDU().deleteKey ("MCRX_PP1_RESUME" );    
    cout << "Resuming postprocessing with camera " << i << endl;
  }
  catch (CCfits::HDU::NoSuchKeyword& h) {}   

  for (; i < n_cameras; ++i) {
    const int j=i;
    std::ostringstream ost;
    ost << j;
    const string camera_numstr=ost.str();

    // By closing and reopening the output file we don't keep the
    // large HDU's all in memory
    output.reset(new FITS (output_file_name, Write));

    // in case program is killed, write where we were to the keyword
    // so we can resume
    output->pHDU().addKey("MCRX_PP1_RESUME", j,
			  "Postprocessing interrupted at this camera" );


    // since the images are surface brightness, we convert them to
    // (inferred) luminosity by back-calculating the internal flux
    // units by dividing by sb_factr, and then multiplying by
    // 4*pi*d^2, where d is the camera distance in whatever units it
    // was specified. After this, the units are L_lambda units.
    T_float distance_factor;
    {
      ExtHDU& parameters = open_HDU (*output, "CAMERA" + camera_numstr+
				     "-PARAMETERS");
      T_float distance;
      parameters.readKey("cameradist", distance);
      distance_factor = 4*constants::pi*distance*distance;
    }

    array_1 scatter_sed
      (integrate_image(*output,  "CAMERA" + camera_numstr, 
		       distance_factor,
		       "L_lambda_scatter" + camera_numstr,
		       wavelengths,
		       "L_scat" +camera_numstr,
		       units,
		       bind(&Mcrx::terminator, this)
		       ));

    // Check for terminator
    if (terminator ()) break;

    array_1 nonscatter_sed
      (integrate_image(*output,
		       "CAMERA" + camera_numstr + "-NONSCATTER", 
		       distance_factor,
		       "L_lambda_nonscatter" + camera_numstr,
		       wavelengths,
		       "L_ns" +camera_numstr,
		       units,
		       bind(&Mcrx::terminator, this)
		       ));

    // Check for terminator
    if (terminator ()) break;

    array_1 ir_sed
      (integrate_image(*output,
		       "CAMERA" + camera_numstr + "-IR", 
		       distance_factor,
		       "L_lambda_ir" + camera_numstr,
		       wavelengths,
		       "L_ir" +camera_numstr,
		       units,
		       bind(&Mcrx::terminator, this)
		       ));

    {
      const string column_name = "L_lambda_out" + camera_numstr;
      array_1 out_sed(scatter_sed+ir_sed);
      const T_float L_bol = integrate_quantity (out_sed, wavelengths, true);
      // now write the total (scatter+ir) sed to its own column 
      ExtHDU& iq_HDU = open_HDU (*output, "INTEGRATED_QUANTITIES"); 
      iq_HDU.addKey("L_out" +camera_numstr, L_bol,
		    "[" + units.get("luminosity") +
		    "] Total bolometric luminosity in camera "+
		    camera_numstr+", including dust emission");
      try {
	iq_HDU.column(column_name);
      }
      catch (Table::NoSuchColumn&) {
	iq_HDU.addColumn(Tdouble, column_name, 1,
			 units.get("L_lambda"));
      }
      write (iq_HDU.column(column_name), out_sed, 1);
    }
    // Check for terminator
    if (terminator ()) break;

    if (dump_sed) {
      for (array_1::iterator i = wavelengths.begin(),
	     j = nonscatter_sed.begin(),
	     l = scatter_sed.begin();
	   i != wavelengths.end(); ++i, ++j, ++l) {
	dump  <<*i << '\t' <<*j  << '\t' <<*l 
	      <<'\n';
      } 
      dump << "\n\n";
    }


  } // for camera

  if (terminator ()) {
    // we were terminated in the camera loop.  The resume keyword is
    // already written, so we don't need to do this.
    cout << "Postprocessing terminated" << endl;
  }
  else {
    // we are done, write keyword to inform of this fact
    output->pHDU().addKey("MCRX_PP1", true,
			  "Postprocessing stage 1 complete" );
    // and delete the resume keyword since we are complete
    output->pHDU().deleteKey ("MCRX_PP1_RESUME" );    
    cout << "Postprocessing stage 1 complete" << endl;
  }
  
  return terminator ();
}  

/** Calculates the total absorbed luminosity in the grid. For
    polychromatic runs, this just consists of coadding the deposition
    spectra. */
mcrx::array_1 mcrx::Mcrx::calculate_L_absorbed(FITS& output)
{
  // if we are running poly this step simply consists of coadding the
  // depositions spectra, so first check that
  bool poly = false;
  ExtHDU& scattering_lambdas_hdu = open_HDU (output, "SCATTERING_LAMBDAS");
  try {
    scattering_lambdas_hdu.readKey("POLY", poly);
  }
  catch (...) {}

  cout << "Calculating absorbed luminosity" << endl;

  ExtHDU& iq_HDU = open_HDU (output, "INTEGRATED_QUANTITIES");
  Column& cl = iq_HDU.column("lambda_intensity" );
  array_1 wavelengths;
  read (cl, wavelengths);
  int n_wavelengths;
  iq_HDU.readKey("NLAMBDA_INTENSITY", n_wavelengths);
  wavelengths.resizeAndPreserve(n_wavelengths);

  array_1 deposition_sed;//(wavelengths.size()); 

  if(poly) {
    cout << "  Reading DEPOSITION" << endl;
    // read deposition image (spectra for all cells)
    try {
      ExtHDU& dep_HDU = open_HDU (output, "DEPOSITION");
      array_2 dep_img;
      read(dep_HDU, dep_img);
      dep_img = pow(10., dep_img);
      // sum over cells
      cout << "  Integrating image" << endl;
      deposition_sed.resize(dep_img.extent(secondDim));
      deposition_sed = sum(dep_img(tensor::j, tensor::i), tensor::j);
    }
    catch (CCfits::FITS::NoSuchHDU&) {
      cout << "No DEPOSITION data found" << endl;
      deposition_sed.resize(wavelengths.size());
      deposition_sed=0;
    }
  }
  else {
    // not poly, we have to do it the old, hard way!

    // read scattering_lambdas stuff
    ExtHDU& lambda_hdu = open_HDU (output, "LAMBDA");
    Column& ce = scattering_lambdas_hdu.column(sllec );
    vector<long> entry;
    vector<bool> line;
    scattering_lambdas_hdu.makeThisCurrent();
    ce.read(entry, 2, ce.rows() );
    scattering_lambdas_hdu.column("line" ).read(line, 2, ce.rows ());

    vector<T_float> scattering_lambdas, continuum_lambdas;
    for (int i = 0 ; i < entry.size() ;++i) {
      scattering_lambdas.push_back( log10 (wavelengths (entry [i])));
      if (!line [i])
	continuum_lambdas.push_back(scattering_lambdas.back());
    }

    // calculate deposition fraction to full resolution
    array_1 deposition_fraction (wavelengths.shape());      
    calc_deposition_fraction(output,
			     continuum_lambdas, 
			     wavelengths, 
			     entry,
			     line, 
			     deposition_fraction);
    assert (all (deposition_fraction == deposition_fraction));
    assert (all (deposition_fraction >= 0));

    Column& cll = iq_HDU.column("L_lambda" );
    array_1 L_lambda_total;
    read (cll, L_lambda_total);
    
    // calculate deposition vector to full wavelength resolution
    deposition_sed.resize(wavelengths.shape());
    deposition_sed = deposition_fraction*L_lambda_total;
  }

  
  // expand vector, if necessary
  assert(deposition_sed.size()==wavelengths.size());
  //deposition_sed.reference(expand(deposition_sed, 
  //			  wavelengths.size()/deposition_sed.size()));

  // now write result
  iq_HDU.makeThisCurrent(); 
  // Necessary, otherwise it is written to
  // the deposition HDU (obviously a bug in CCfits) 
  write(iq_HDU.column("L_lambda_absorbed"), deposition_sed, 1);
  const T_float L_bol_absorbed = integrate_quantity (deposition_sed,
						     wavelengths, true);
  iq_HDU.addKey("L_bol_absorbed", L_bol_absorbed, "[" + 
		units.get("luminosity") +
		"] Bolometric luminosity absorbed in grid");
  cout << "Absorbed bolometric luminosity: " << L_bol_absorbed << " " <<
    units.get("luminosity") << endl;

  return deposition_sed;
}

